import sys
from pathlib import Path
from collections import OrderedDict

import numpy as np
import onnxruntime as ort
from ais_bench.infer.interface import InferSession


class Comparator:

    def __init__(self, model_path1, model_path2, shape_info=None, om_mode=None):
        self.om_mode = om_mode
        self.helper1 = self.create_inference_helper(model_path1, shape_info)
        self.helper2 = self.create_inference_helper(model_path2, shape_info)

        msg = f'{self.helper1.input_names} != {self.helper2.input_names}'
        assert self.helper1.input_names == self.helper2.input_names, msg

        self.fix_input_shape(shape_info)

    def fix_input_shape(self, shape_info=None):
        shape_dict = {}
        if shape_info:
            for text in shape_info.split(';'):
                name, shape = text.rsplit(':', 1)
                shape = [c if c == 'x' else int(c) for c in shape.strip().split(',')]
                shape_dict[name.strip()] = shape

        input_names = self.helper1.input_names
        shapes1 = self.helper1.input_shapes
        shapes2 = self.helper2.input_shapes
        assert len(input_names) == len(shapes1) == len(shapes2)

        fixed_shapes = []
        for name, s1, s2 in zip(input_names, shapes1, shapes2):
            s3 = shape_dict[name] if name in shape_dict else None
            assert len(s1) == len(s2)
            if s3: 
                assert len(s3) == len(s1), f'Must specify a valid input shape.'
            shape = []
            for i, (x1, x2) in enumerate(zip(s1, s2)):
                size = None
                x3 = None if s3 is None else s3[i]
                if isinstance(x1, int) and x1 > 0:
                    size = x1
                if isinstance(x2, int) and x2 > 0:
                    size = x2
                if isinstance(x3, int) and x3 > 0:
                    size = x3
                shape.append(size)
            fixed_shapes.append(shape)

        self.helper1.fix_input_shape(fixed_shapes)
        self.helper2.fix_input_shape(fixed_shapes)

    def create_inference_helper(self, model_path, shape_info):
        suffix = Path(model_path).suffix
        if suffix == '.onnx':
            return OnnxInferenceHelper(model_path, shape_info)
        elif suffix == '.om':
            return OmInferenceHelper(model_path, shape_info, self.om_mode)
        else:
            raise Exception(f'Unknown model type: {suffix}')

    def compare(self, data_paths=None):
        if data_paths:
            data = OrderedDict()
            for i, data_path in enumerate(data_paths):
                if not Path(data_path).is_file():
                    raise Exception(f'File not exists: {data_path}')
                d = np.load(data_path)
                assert list(d.shape) == self.helper1.input_shapes[i], \
                    f'{d.shape}, {self.helper1.input_shapes[i]}'
                data[self.helper1.input_names[i]] = d
        else:
            data = self.helper1.prepare_data()
        outputs1 = self.helper1.inference(data)
        outputs2 = self.helper2.inference(data)
        self.compare_outputs(outputs1, outputs2)

    def cosine_similarity(self, a, b):
        a_ = a.flatten().astype(np.float64)
        b_ = b.flatten().astype(np.float64)
        cos = np.dot(a_, b_) / (np.linalg.norm(a_) * np.linalg.norm(b_))
        return cos

    def difference(self, a, b):
        a_ = a.flatten().astype(np.float64)
        b_ = b.flatten().astype(np.float64)
        diff = np.abs(a_ - b_)
        max_diff = np.max(diff)
        avg_diff = np.mean(diff)
        sum_diff = np.sum(diff)
        return max_diff, avg_diff, sum_diff

    def compare_outputs(self, outputs1, outputs2):
        for i, (o1, o2) in enumerate(zip(outputs1, outputs2)):
            print(f"compare output {i}", ">" * 50)
            cos = self.cosine_similarity(o1, o2)
            max_diff, avg_diff, sum_diff = self.difference(o1, o2)
            all_close = np.allclose(o1, o2, rtol=1e-05, atol=1e-08, equal_nan=True)
    
            print("cosine similarity:", cos)
            print("max(abs(a - b)):", max_diff)
            print("avg(abs(a - b)):", avg_diff)
            print("sum(abs(a - b)):", sum_diff)
            print("all close:", all_close)
            print("<" * 67)


class BaseInferenceHelper:
    def __init__(self):
        self.model_type = None
        self.session = None
        self.input_names = None
        self.input_shapes = None
        self.input_dtypes = None

    def fix_input_shape(self, fixed_shapes=None):
        assert len(fixed_shapes) == len(self.input_shapes)
        for s1, s2 in zip(self.input_shapes, fixed_shapes):
            assert len(s1) == len(s2), f'Must specify a valid input shape.'
            for i, (x1, x2) in enumerate(zip(s1, s2)):
                if isinstance(x1, int) and x1 > 0:
                    assert x2 == 'x' or x2 == x1, f'input shape conflict: {s1} <--> {s2}'
                else:
                    assert isinstance(x2, int) and x2 > 0, f'Must specify a valid input shape.'
                    s1[i] = x2

    def dtype_convert(self, data_type):
        err_msg = f'Please add the convert rule for dtype: {data_type}'
        raise NotImplementedError(err_msg)

    def prepare_data(self):
        input_feed = OrderedDict()
        for i, name in enumerate(self.input_names):
            shape = self.input_shapes[i]
            dtype = self.input_dtypes[i]
            print(f'input shape: {shape}    input dtype: {dtype}')
            data = np.random.randn(*shape).astype(dtype)
            input_feed[name] = data
        return input_feed

    def inference(self, input_feed=None):
        if input_feed is None:
            input_feed = self.prepare_data()

        if self.model_type == 'om':
            input_feed = list(input_feed.values())
            return self.session.infer(input_feed, self.mode)
        elif self.model_type == 'onnx':
            return self.session.run(self.output_names, input_feed)
        else:
            raise Exception(f'Unknown model type: {self.model_type}')


class OnnxInferenceHelper(BaseInferenceHelper):
    def __init__(self, onnx_path, shape_info=None):
        super(OnnxInferenceHelper, self).__init__()
        self.model_type = 'onnx'
        self.session = ort.InferenceSession(onnx_path)
        input_info = self.session.get_inputs()
        self.input_names = [i.name for i in input_info]
        self.input_shapes = [i.shape for i in input_info]
        self.input_dtypes = [self.dtype_convert(i.type) for i in input_info]
        self.output_names = [o.name for o in self.session.get_outputs()]

    def dtype_convert(self, type_str):
        if type_str == 'tensor(float)':
            return 'float32'
        if type_str == 'tensor(int64)':
            return 'int64'
        err_msg = f'Please add the convert rule for dtype: {type_str}'
        raise NotImplementedError(err_msg)


class OmInferenceHelper(BaseInferenceHelper):
    def __init__(self, om_path, shape_info=None, om_mode=None):
        super(OmInferenceHelper, self).__init__()
        self.model_type = 'om'
        self.mode = om_mode
        device_id = 0
        self.session = InferSession(device_id, om_path)
        input_info = self.session.get_inputs()
        self.input_names = [i.name for i in input_info]
        self.input_shapes = [i.shape for i in input_info]
        self.input_dtypes = [i.datatype.name for i in input_info]
        self.output_names = [o.name for o in self.session.get_outputs()]


def main():
    from argparse import ArgumentParser
    parser = ArgumentParser()
    parser.add_argument('model_path1', type=str, 
                        help='path to the first model(ONNX or OM).')
    parser.add_argument('model_path2', type=str, 
                        help='path to the second model(ONNX or OM).')
    parser.add_argument('-s', '--shapes', type=str, default=None,
                        help='specify value for dynamic dimensions, eg: '
                        '"input1:4,x;input2:4,3,x,x", "x" means it is a fixed dimension.')
    parser.add_argument('-d', '--data_paths', type=str, nargs='*', 
                        help='path to test data(npy files).')
    parser.add_argument('-m', '--om_mode', type=str, default='static', 
                        choices=['static', 'dymdims', 'dymshape'],
                        help='mode of OM model, related to the dynamic '
                        'input type of the model')
    args = parser.parse_args()

    comparator = Comparator(args.model_path1, args.model_path2, 
                            shape_info=args.shapes, om_mode=args.om_mode)
    comparator.compare(data_paths=args.data_paths)


if __name__ == "__main__":
    main()
